package com.study.interceptor;

import com.study.utils.IpAddressUtil;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.data.redis.core.ZSetOperations;
import org.springframework.stereotype.Component;
import org.springframework.web.servlet.HandlerInterceptor;

import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import java.util.concurrent.TimeUnit;

/**
 * 限制同一时间段内来自同一个IP地址对同一个方法的的请求数。
 */
@Component
@Slf4j
public class RateLimitInterceptor implements HandlerInterceptor {

    private static final int MAX_REQUESTS = 30; // 同一时间段内允许的最大请求数
    private static final long TIME_PERIOD = 60 * 1000; // 时间段，单位为毫秒
    @Autowired
    private RedisTemplate redisTemplate;

    @Override
    public boolean preHandle(HttpServletRequest request, HttpServletResponse response, Object handler) {
        // request info
        String ipAddress = IpAddressUtil.getIpAddr(request);
        String requestPath = request.getRequestURI();
        try {
            // key
            String key = "req_limit:req_limit_".concat(requestPath).concat(ipAddress);
            // 获取redis操作对象
            ZSetOperations zSetOperations = redisTemplate.opsForZSet();

            // 添加当前时间戳
            long currentMs = System.currentTimeMillis();
            zSetOperations.add(key, currentMs, currentMs);
            // 设置代码用户的过期时间
            redisTemplate.expire(key, TIME_PERIOD, TimeUnit.MILLISECONDS);
            // 删除当前窗口外的值
            zSetOperations.removeRangeByScore(key, 0, currentMs - TIME_PERIOD);

            // 检查所有可用计数
            Long count = zSetOperations.zCard(key);
            if (count > MAX_REQUESTS) {
                response.setStatus(429); //设置响应状态码
                response.getWriter().write("Too many requests from this IP address");
                log.error("接口拦截：{} 请求超过限制频率【{}次/{}s】,IP为{}", requestPath, MAX_REQUESTS, TIME_PERIOD / 1000, ipAddress);
                return false;
            }
            // 执行用户请求
            return true;
        } catch (Exception e) {
            log.error("IP{}请求处理'{}'异常：{}", ipAddress, requestPath, e);
            response.setStatus(404); //设置响应状态码
            return false;
        }
    }
}
